from annoy import AnnoyIndex
import numpy as np
import os

# 1. 数据生成
np.random.seed(42)
num_vectors = 10000  # 模拟1万条向量
dimension = 50  # 每条向量50维
data = np.random.rand(num_vectors, dimension)  # 随机生成数据
query_point = np.random.rand(1, dimension)  # 查询点

# 2. 构建Annoy索引
num_trees = 20  # 设置随机投影树的数量

index_path = "annoy_index.ann"  # 索引文件路径
annoy_index = AnnoyIndex(dimension, metric='euclidean')

# 添加数据索引
for i in range(num_vectors):
    annoy_index.add_item(i, data[i])

annoy_index.build(num_trees)  # 构建索引
annoy_index.save(index_path)  # 持久化索引到磁盘

# 3. 加载索引并查询
loaded_index = AnnoyIndex(dimension, metric='euclidean')
loaded_index.load(index_path)  # 从磁盘加载索引

# 查询最近邻
top_k = 5
search_k = 1000  # 控制搜索的最大树数量，优化查询速度
results = loaded_index.get_nns_by_vector(
    query_point[0], top_k, search_k=search_k, include_distances=True)

# 4. 输出查询结果
print("查询点:", query_point[0])
print("\n最近邻查询结果:")
for idx, (neighbor_idx, distance) in enumerate(zip(results[0], results[1])):
    print(f"结果 {idx + 1}: 索引 {neighbor_idx}, 距离 {distance:.4f}")

# 5. 索引性能分析
print("\n性能优化分析:")
print(f"索引中随机投影树的数量: {num_trees}")
print(f"搜索时使用的树数量: {search_k}")
print(f"索引文件大小: {os.path.getsize(index_path) / (1024 * 1024):.2f} MB")

# 删除持久化文件（清理环境）
os.remove(index_path)